import numpy as np
import gradio as gr
from PIL import Image
from pinecone import Pinecone
import os

api_key = "ecb9978d-a69b-458c-aa23-37d03b1407f7"
environment = "us-east-1"
pc = Pinecone(api_key=api_key)
index_name = "quickstart"


index = pc.Index(index_name)

def preprocess_image(input_data):
    if isinstance(input_data, dict):
        image_data = input_data.get('composite', None)
        if image_data is None:
            image_data = input_data.get('image', None)
    else:
        image_data = input_data

    if image_data is None or np.all(image_data == 0):
        print("Empty or invalid image data received.")
        return None

    pil_image = Image.fromarray(image_data).convert('L')
    pil_image = pil_image.resize((8, 8), Image.LANCZOS)
    image_array = np.array(pil_image).astype(float)
    image_array = (image_array / 255.0 * 16.0)  
    return image_array.flatten()

def knn_predict(query_vector, k=11):
    results = index.query(vector = query_vector.tolist(), top_k=k, include_metadata=True)
    votes = [match['metadata']['label'] for match in results['matches']]
    return max(set(votes), key=votes.count)

def predict_digit(image):
    if image is None:
        return {}, ""
    
    processed_image = preprocess_image(image)
    
    if processed_image is None:
        return {}, "Error: Invalid image"

    prediction = knn_predict(processed_image)
    

    k = 11
    results = index.query(vector=processed_image.tolist(), top_k=k, include_metadata=True)
    votes = [match['metadata']['label'] for match in results['matches']]
    vote_counts = {i: votes.count(i) for i in range(10)}
    total_votes = sum(vote_counts.values())
    probabilities = {str(i): count / total_votes for i, count in vote_counts.items()}
    
    return probabilities, str(prediction)

# Create Gradio interface
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="Draw a digit here",mode="L"),
    outputs=[gr.Label(num_top_classes=3), gr.Textbox(label="Prediction")],
    live=False,
    title="Handwritten Digit Recognition using Pinecone KNN",
    description="Pinecone KNN model -- draw a digit and see the prediction!",
    allow_flagging="never",
)

# Launch Gradio interface
iface.launch(share=True)